import os
import shutil
from pathlib import Path
import streamlit as st
import pandas as pd
from obspy import read
import tkinter as tk
from tkinter import filedialog

from eps2_to_sac import convert_eps2_to_sac, convert_eps2_directory_to_sac
from sac_sync import batch_synchronize_sac
from sac_to_segy import (
    convert_single_sac_to_output,
    convert_directory_sac_to_segy,
    combine_sacs_to_single_segy,
    list_sac_files,
)
from segy_to_seg2 import convert_segy_to_seg2, convert_directory_segy_to_seg2


st.set_page_config(page_title="地震数据转换工具", layout="wide")

st.title("地震数据转换工具")
st.caption("支持EPS2到SEGY的完整转换流程：EPS2→SAC→时间同步→SEGY")


def create_workflow_dir(base_dir: Path, workflow_name: str):
    """创建工作流目录结构"""
    workflow_dir = base_dir / workflow_name
    sac_dir = workflow_dir / "sac"
    sync_dir = workflow_dir / "sync"
    segy_dir = workflow_dir / "segy"
    
    # 创建目录
    sac_dir.mkdir(parents=True, exist_ok=True)
    sync_dir.mkdir(parents=True, exist_ok=True)
    segy_dir.mkdir(parents=True, exist_ok=True)
    
    return workflow_dir, sac_dir, sync_dir, segy_dir


def list_eps2_files(directory: Path, recursive: bool = False) -> list[Path]:
    """列出目录中的所有EPS2文件，通过读取DATAFILE.LST文件获取实际存在的文件列表"""
    
    def get_eps2_files_from_dir(dir_path: Path) -> list[Path]:
        """从单个目录获取EPS2文件列表"""
        eps2_files = []
        
        # 检查是否存在DATAFILE.LST文件
        datalist_file = dir_path / "DATAFILE.LST"
        if datalist_file.exists():
            try:
                with open(datalist_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        
                        # 解析DATAFILE.LST文件格式：序号 日期 时间 1:文件名.后缀 采样率
                        parts = line.split()
                        if len(parts) >= 5:
                            # 查找包含冒号的文件名部分（如1:66BC75C6.012）
                            filename_part = None
                            for part in parts:
                                if ':' in part and part.startswith('1:'):
                                    filename_part = part
                                    break
                            
                            if filename_part and ':' in filename_part:
                                # 提取实际文件名（如66BC75C6.012）
                                actual_filename = filename_part.split(':')[1]
                                
                                # 检查文件是否存在
                                file_path = dir_path / actual_filename
                                if file_path.exists():
                                    eps2_files.append(file_path)
            except Exception as e:
                print(f"读取DATAFILE.LST文件失败 {datalist_file}: {e}")
                # 如果读取失败，回退到直接查找三位数字后缀的文件
                pattern_files = list(dir_path.glob("*.???"))
                for file_path in pattern_files:
                    if file_path.suffix and len(file_path.suffix) == 4 and file_path.suffix[1:].isdigit():
                        eps2_files.append(file_path)
        else:
            # 如果没有DATAFILE.LST文件，直接查找三位数字后缀的文件
            pattern_files = list(dir_path.glob("*.???"))
            for file_path in pattern_files:
                if file_path.suffix and len(file_path.suffix) == 4 and file_path.suffix[1:].isdigit():
                    eps2_files.append(file_path)
        
        return eps2_files
    
    # 根据recursive参数决定搜索方式
    if recursive:
        # 递归搜索所有子目录
        all_eps2_files = []
        for subdir in directory.rglob("*"):
            if subdir.is_dir():
                all_eps2_files.extend(get_eps2_files_from_dir(subdir))
        return all_eps2_files
    else:
        # 只搜索当前目录
        return get_eps2_files_from_dir(directory)


def estimate_segy_parts(npts: int, max_samples: int = 32767) -> int:
    """按 SEGY 写出限制估算需要的分割份数（时间窗口）。"""
    if npts <= 0:
        return 0
    return (int(npts) + max_samples - 1) // max_samples


def build_combine_mapping(sac_files: list[Path]) -> dict:
    """构建合并映射信息：按读取顺序列出每个输入 trace 的采样信息与分块数。"""
    dt_us_ref = None
    npts_ref = None
    inconsistent: list[Path] = []
    over_limit: list[Path] = []
    entries = []
    for sac_path in sac_files:
        st_obj = read(str(sac_path))
        for ti, tr in enumerate(st_obj, start=1):
            dt_us = int(round(tr.stats.delta * 1_000_000))
            npts = int(tr.stats.npts)
            if dt_us_ref is None:
                dt_us_ref = dt_us
                npts_ref = npts
            else:
                if dt_us != dt_us_ref or npts != npts_ref:
                    inconsistent.append(sac_path)
            chunks = estimate_segy_parts(npts)
            if npts > 32767:
                over_limit.append(sac_path)
            entries.append({
                "sac_path": sac_path,
                "trace_in_sac": ti,
                "dt_us": dt_us,
                "npts": npts,
                "chunks": chunks,
            })
    return {
        "entries": entries,
        "dt_us_ref": dt_us_ref,
        "npts_ref": npts_ref,
        "inconsistent": inconsistent,
        "over_limit": over_limit,
    }


def select_directory() -> str:
    """打开目录选择对话框"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    root.attributes('-topmost', True)  # 置顶窗口
    
    directory = filedialog.askdirectory(title="选择目录")
    root.destroy()
    
    return directory if directory else ""


with st.sidebar:
    st.header("工作流选项")
    workflow_mode = st.radio(
        "选择工作流模式",
        options=["完整工作流(EPS2→SEGY)", "分步模式", "EPS2→SAC", "SAC时间同步", "SAC→SEGY", "验证完整流程模式"],
        index=0
    )
    
    if workflow_mode == "完整工作流(EPS2→SEGY)" or workflow_mode == "分步模式":
        workflow_name = st.text_input("工作流名称", value="workflow_test")
        base_output_dir = st.text_input("输出目录", value=str(Path.cwd()))
        
        if workflow_mode == "完整工作流(EPS2→SEGY)":
            st.info("完整工作流将执行：EPS2→SAC→时间同步→SEGY")
            no_sync = st.checkbox("跳过时间同步", value=False)
            if not no_sync:
                target_sampling_rate = st.number_input("目标采样率 (Hz)", value=100.0, min_value=0.0, format="%.2f", help="指定重采样时的目标采样率。如果设为0，则自动使用最低采样率。")
            
            no_combine = st.checkbox("不合并为单个SEGY文件", value=False)
            recursive = st.checkbox("递归处理子目录", value=True)
            strict_one = st.checkbox("严格模式：一个SAC文件对应一个SEGY道", value=False)
            split_output_files = st.checkbox("按样本窗口分割为多个SEGY文件", value=False)
            merge_z_only = st.checkbox("只合并Z分量 (文件名以 _Z.sac 结尾)", value=False)
            output_seg2 = st.checkbox("同时输出SEG2格式", value=False)
            
            with st.expander("参数说明"):
                st.markdown(
                    "- `跳过时间同步`：不执行SAC文件时间同步，直接转换。\n"
                    "- `目标采样率`：指定同步时所有文件重采样的目标频率。\n"
                    "- `不合并为单个SEGY文件`：每个SAC文件转换为单独的SEGY文件。\n"
                    "- `递归处理子目录`：处理输入目录及其所有子目录中的EPS2文件。\n"
                    "- `严格模式`：一个SAC文件对应一个SEGY道，超过32767样本将报错。\n"
                    "- `按样本窗口分割`：将长SAC文件分割为多个SEGY文件。\n"
                    "- `只合并Z分量`：仅合并Z分量文件。\n"
                    "- `同时输出SEG2格式`：在生成SEGY文件的同时，自动生成同名的SEG2文件。"
                )
        
        elif workflow_mode == "分步模式":
            st.info("分步模式允许您逐步执行每个转换阶段")
            step1 = st.checkbox("步骤1: EPS2→SAC", value=True)
            step2 = st.checkbox("步骤2: SAC时间同步", value=True)
            if step2:
                target_sampling_rate_step = st.number_input("目标采样率 (Hz)", value=100.0, min_value=0.0, format="%.2f", key="sampling_rate_step", help="指定重采样时的目标采样率。如果设为0，则自动使用最低采样率。")

            step3 = st.checkbox("步骤3: SAC→SEGY", value=True)
            
            if step3:
                segy_mode = st.radio("SEGY转换模式", ["合并为单个SEGY", "每个SAC单独转换"])
                if segy_mode == "合并为单个SEGY":
                    strict_one = st.checkbox("严格模式：一个SAC文件对应一个SEGY道", value=False)
                    split_output_files = st.checkbox("按样本窗口分割为多个SEGY文件", value=False)
                    merge_z_only_step = st.checkbox("只合并Z分量 (文件名以 _Z.sac 结尾)", value=False, key="z_only_step")
                else:
                    strict_one = False
                output_seg2 = st.checkbox("同时输出SEG2格式", value=False)
                split_output_files = False
                merge_z_only_step = False
    
    elif workflow_mode == "EPS2→SAC":
        col1, col2 = st.columns([3, 1])
        with col1:
            eps2_input_dir = st.text_input("EPS2输入目录", value=str(Path.cwd() / "EPS2data"), key="eps2_input_dir_eps2sac")
        with col2:
            if st.button("浏览目录", key="browse_eps2_eps2sac"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.eps2_input_dir_eps2sac = selected_dir
        
        col3, col4 = st.columns([3, 1])
        with col3:
            sac_output_dir = st.text_input("SAC输出目录", value=str(Path.cwd() / "sac_out"), key="sac_output_dir_eps2sac")
        with col4:
            if st.button("浏览目录", key="browse_sac_eps2sac"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.sac_output_dir_eps2sac = selected_dir
        
        # 从session_state获取目录值
        if 'eps2_input_dir_eps2sac' in st.session_state:
            eps2_input_dir = st.session_state.eps2_input_dir_eps2sac
        if 'sac_output_dir_eps2sac' in st.session_state:
            sac_output_dir = st.session_state.sac_output_dir_eps2sac
        
        recursive = st.checkbox("递归处理子目录", value=True)
    
    elif workflow_mode == "SAC时间同步":
        col1, col2 = st.columns([3, 1])
        with col1:
            sac_input_dir = st.text_input("SAC输入目录", value=str(Path.cwd() / "sacdata"), key="sac_input_dir_sync")
        with col2:
            if st.button("浏览目录", key="browse_sac_input_sync"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.sac_input_dir_sync = selected_dir
        
        col3, col4 = st.columns([3, 1])
        with col3:
            sac_sync_output_dir = st.text_input("同步后SAC输出目录", value=str(Path.cwd() / "sac_sync"), key="sac_sync_output_dir_sync")
        with col4:
            if st.button("浏览目录", key="browse_sac_output_sync"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.sac_sync_output_dir_sync = selected_dir
        
        # 从session_state获取目录值
        if 'sac_input_dir_sync' in st.session_state:
            sac_input_dir = st.session_state.sac_input_dir_sync
        if 'sac_sync_output_dir_sync' in st.session_state:
            sac_sync_output_dir = st.session_state.sac_sync_output_dir_sync
        
        target_sampling_rate_sync = st.number_input("目标采样率 (Hz)", value=100.0, min_value=0.0, format="%.2f", key="sampling_rate_sync", help="指定重采样时的目标采样率。如果设为0，则自动使用最低采样率。")
        
        component_groups = st.text_area(
            "分量组配置（每行一组，用逗号分隔）",
            value="X,Y,Z\nX,Y\nX,Z\nY,Z"
        )
    
    elif workflow_mode == "SAC→SEGY":
        col1, col2 = st.columns([3, 1])
        with col1:
            sac_segy_input_dir = st.text_input("SAC输入目录", value=str(Path.cwd() / "sacdata"), key="sac_segy_input_dir")
        with col2:
            if st.button("浏览目录", key="browse_sac_segy_input"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.sac_segy_input_dir = selected_dir
        
        col3, col4 = st.columns([3, 1])
        with col3:
            segy_output_dir = st.text_input("SEGY输出目录", value=str(Path.cwd() / "segy_out"), key="segy_output_dir")
        with col4:
            if st.button("浏览目录", key="browse_segy_output"):
                selected_dir = select_directory()
                if selected_dir:
                    st.session_state.segy_output_dir = selected_dir
        
        # 从session_state获取目录值
        if 'sac_segy_input_dir' in st.session_state:
            sac_segy_input_dir = st.session_state.sac_segy_input_dir
        if 'segy_output_dir' in st.session_state:
            segy_output_dir = st.session_state.segy_output_dir
        
        segy_mode = st.radio("SEGY转换模式", ["合并为单个SEGY", "每个SAC单独转换"])
        if segy_mode == "合并为单个SEGY":
            strict_one = st.checkbox("严格模式：一个SAC文件对应一个SEGY道", value=False)
            split_output_files = st.checkbox("按样本窗口分割为多个SEGY文件", value=False)
            merge_z_only_sac = st.checkbox("只合并Z分量 (文件名以 _Z.sac 结尾)", value=False, key="z_only_sac")
        else:
            strict_one = False
            split_output_files = False
            merge_z_only_sac = False


# 完整工作流模式
if workflow_mode == "完整工作流(EPS2→SEGY)":
    st.header("完整工作流：EPS2→SAC→时间同步→SEGY")
    
    col1, col2 = st.columns([3, 1])
    with col1:
        eps2_input_dir = st.text_input("EPS2文件目录", value="D:\20251030试验数据\20251030试验数据", key="eps2_input_dir_full")
    with col2:
        if st.button("浏览目录", key="browse_eps2_full"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.eps2_input_dir_full = selected_dir
    
    col3, col4 = st.columns([3, 1])
    with col3:
        base_output_dir = st.text_input("输出目录", value=str(Path.cwd()), key="base_output_dir_full")
    with col4:
        if st.button("浏览目录", key="browse_output_full"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.base_output_dir_full = selected_dir
    
    # 从session_state获取目录值
    if 'eps2_input_dir_full' in st.session_state:
        eps2_input_dir = st.session_state.eps2_input_dir_full
    if 'base_output_dir_full' in st.session_state:
        base_output_dir = st.session_state.base_output_dir_full
    
    run_workflow = st.button("执行完整工作流")
    
    if run_workflow:
        try:
            base_dir = Path(base_output_dir)
            workflow_dir, sac_dir, sync_dir, segy_dir = create_workflow_dir(base_dir, workflow_name)
            
            # 转换EPS2到SAC
            sac_files = []
            with st.spinner("步骤1: 转换EPS2到SAC..."):
                eps2_dir = Path(eps2_input_dir)
                if not eps2_dir.is_dir():
                    st.error(f"EPS2目录不存在: {eps2_dir}")
                else:
                    eps2_files = list_eps2_files(eps2_dir, recursive=recursive)
                    if not eps2_files:
                        st.error(f"在目录中未找到EPS2文件: {eps2_dir}")
                    else:
                        st.info(f"找到 {len(eps2_files)} 个EPS2文件，开始转换...")
                        for eps2_file in eps2_files:
                            try:
                                files = convert_eps2_to_sac(eps2_file, sac_dir)
                                sac_files.extend(files)
                                st.success(f"✓ 成功转换 {eps2_file.name} -> {len(files)} 个SAC文件")
                            except Exception as e:
                                st.error(f"✗ 转换失败 {eps2_file.name}: {e}")
            
            if sac_files:
                st.success(f"步骤1完成: 成功转换 {len(sac_files)} 个SAC文件为SAC格式")
            
            # SAC文件时间同步
            synced_sac_files = []
            if not no_sync:
                with st.spinner("步骤2: SAC文件时间同步..."):
                    if sac_files:
                        synced_sac_files = batch_synchronize_sac(
                            sac_files, 
                            sync_dir,
                            target_sampling_rate=target_sampling_rate
                        )
                        st.success(f"步骤2完成: 成功同步 {len(synced_sac_files)} 个SAC文件")
                    else:
                        st.warning("没有需要同步的SAC文件。")
            
            # 转换SAC到SEGY
            with st.spinner("步骤3: 转换SAC到SEGY..."):
                segy_files = []
                trace_mapping = []
                
                # 根据是否同步，选择处理的SAC文件列表
                sac_files_for_segy = synced_sac_files if not no_sync and synced_sac_files else sac_files
                source_dir_for_segy = sync_dir if not no_sync and synced_sac_files else sac_dir

                if not sac_files_for_segy:
                    st.warning("没有找到SAC文件用于SEGY转换。")
                else:
                    if no_combine:
                        # 每个SAC文件单独转换
                        segy_files = convert_directory_sac_to_segy(
                            source_dir_for_segy, 
                            segy_dir, 
                            recursive=False, # 因为文件都在一个目录里
                            overwrite=True
                        )
                    else:
                        # 合并为单个SEGY文件
                        st.info(f"开始合并 {len(sac_files_for_segy)} 个SAC文件...")
                        
                        # 如果用户选择只合并Z分量
                        if merge_z_only:
                            sac_files_to_process = [f for f in sac_files_for_segy if Path(f).stem.upper().endswith('_Z')]
                            st.info(f"Z-only模式：找到 {len(sac_files_to_process)} 个Z分量文件进行合并。")
                        else:
                            sac_files_to_process = sac_files_for_segy

                        if sac_files_to_process:
                            output_filename = f"combined_{workflow_name}.segy"
                            if merge_z_only:
                                output_filename = f"combined_{workflow_name}_Z.segy"
                            output_file = segy_dir / output_filename

                            try:
                                trace_mapping = combine_sacs_to_single_segy(
                                    sac_files_to_process,
                                    output_file,
                                    overwrite=True,
                                    strict_one_trace_per_sac=strict_one,
                                    split_output_files=split_output_files,
                                    merge_z_only=False
                                )

                                if split_output_files:
                                    base_pattern = f"{output_file.stem}_part*.segy"
                                    segy_files = [str(p) for p in segy_dir.glob(base_pattern)]
                                    st.success(f"✓ 成功合并并分割为 {len(segy_files)} 个SEGY文件。")
                                else:
                                    segy_files = [str(output_file)]
                                    st.success(f"✓ 成功将 {len(sac_files_to_process)} 个SAC文件合并到 {output_file.name}")

                            except Exception as e:
                                st.error(f"合并SAC文件失败: {e}")
                        else:
                            st.warning("没有找到符合条件的SAC文件进行合并。")
            
            st.success(f"步骤3完成: 成功生成 {len(segy_files)} 个SEGY文件")
            
            # 如果选择了同时输出SEG2格式
            if output_seg2 and segy_files:
                with st.spinner("步骤4: 转换SEGY到SEG2..."):
                    seg2_dir = workflow_dir / "seg2"
                    seg2_dir.mkdir(parents=True, exist_ok=True)
                    seg2_files = []
                    
                    for segy_file in segy_files:
                        try:
                            seg2_file = convert_segy_to_seg2(segy_file, seg2_dir)
                            seg2_files.append(seg2_file)
                            st.success(f"✓ 成功转换 {Path(segy_file).name} -> {Path(seg2_file).name}")
                        except Exception as e:
                            st.error(f"✗ 转换失败 {Path(segy_file).name}: {e}")
                    
                    if seg2_files:
                        st.success(f"步骤4完成: 成功生成 {len(seg2_files)} 个SEG2文件")
            
            # 显示结果
            st.subheader("转换结果")
            st.write(f"工作流目录: {workflow_dir}")
            
            # 显示SAC文件
            st.subheader("生成的SAC文件")
            sac_files_list = list(sac_dir.glob("**/*.sac"))
            if sac_files_list:
                sac_df = pd.DataFrame({"文件名": [f.name for f in sac_files_list], "路径": [str(f) for f in sac_files_list]})
                st.dataframe(sac_df)
            else:
                st.write("未找到SAC文件")

            # 显示同步后的SAC文件
            st.subheader("同步后的SAC文件")
            sync_files_list = list(sync_dir.glob("*.sac"))
            if sync_files_list:
                sync_df = pd.DataFrame({"文件名": [f.name for f in sync_files_list], "路径": [str(f) for f in sync_files_list]})
                st.dataframe(sync_df)
            else:
                st.write("未找到同步后的SAC文件")
            
            # 显示SEGY文件
            st.subheader("生成的SEGY文件")
            if segy_files:
                segy_df = pd.DataFrame({"文件名": [Path(f).name for f in segy_files], "路径": [str(f) for f in segy_files]})
                st.dataframe(segy_df)

                if not no_combine and trace_mapping:
                    st.subheader("文件到道的映射信息")
                    mapping_df = pd.DataFrame({
                        "道号": [i+1 for i in range(len(trace_mapping))],
                        "原始SAC文件名": trace_mapping
                    })
                    st.dataframe(mapping_df)
            else:
                st.write("未找到SEGY文件")
                
            # 显示SEG2文件（如果有）
            if output_seg2 and 'seg2_files' in locals() and seg2_files:
                st.subheader("生成的SEG2文件")
                seg2_df = pd.DataFrame({"文件名": [f.name for f in seg2_files], "路径": [str(f) for f in seg2_files]})
                st.dataframe(seg2_df)
            elif output_seg2:
                st.subheader("生成的SEG2文件")
                st.write("未找到SEG2文件")
            
            # 提供打开输出目录按钮
            if st.button("打开输出目录"):
                try:
                    os.startfile(str(workflow_dir))
                except Exception as e:
                    st.warning(f"无法打开目录：{e}")
        
        except Exception as e:
            st.error(f"工作流执行失败：{e}")


# 分步模式
elif workflow_mode == "分步模式":
    st.header("分步执行转换流程")
    
    col1, col2 = st.columns([3, 1])
    with col1:
        eps2_input_dir = st.text_input("EPS2文件目录", value="EPS2data", key="eps2_input_dir_step")
    with col2:
        if st.button("浏览目录", key="browse_eps2_step"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.eps2_input_dir_step = selected_dir
    
    col3, col4 = st.columns([3, 1])
    with col3:
        base_output_dir = st.text_input("输出目录", value=str(Path.cwd()), key="base_output_dir_step")
    with col4:
        if st.button("浏览目录", key="browse_output_step"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.base_output_dir_step = selected_dir
    
    # 从session_state获取目录值
    if 'eps2_input_dir_step' in st.session_state:
        eps2_input_dir = st.session_state.eps2_input_dir_step
    if 'base_output_dir_step' in st.session_state:
        base_output_dir = st.session_state.base_output_dir_step
    
    run_step1 = st.button("执行步骤1: EPS2→SAC") if step1 else None
    run_step2 = st.button("执行步骤2: SAC时间同步") if step2 else None
    run_step3 = st.button("执行步骤3: SAC→SEGY") if step3 else None
    
    # 步骤1: EPS2到SAC
    if run_step1 and step1:
        try:
            base_dir = Path(base_output_dir)
            workflow_dir, sac_dir, sync_dir, segy_dir = create_workflow_dir(base_dir, workflow_name)
            
            sac_files = []
            with st.spinner("转换EPS2到SAC..."):
                eps2_dir = Path(eps2_input_dir)
                if not eps2_dir.is_dir():
                    st.error(f"EPS2目录不存在: {eps2_dir}")
                else:
                    eps2_files = list_eps2_files(eps2_dir, recursive=True)
                    if not eps2_files:
                        st.error(f"在目录中未找到EPS2文件: {eps2_dir}")
                    else:
                        failed_files = []
                        st.info(f"找到 {len(eps2_files)} 个EPS2文件，开始转换...")
                        for eps2_file in eps2_files:
                            try:
                                files = convert_eps2_to_sac(eps2_file, sac_dir)
                                sac_files.extend(files)
                                st.success(f"✓ 成功转换 {eps2_file.name} -> {len(files)} 个SAC文件")
                            except Exception as e:
                                failed_files.append((eps2_file, str(e)))
                                st.error(f"✗ 转换失败 {eps2_file.name}: {e}")
                        
                        if failed_files:
                            st.warning(f"有 {len(failed_files)} 个EPS2文件转换失败")
            
            if sac_files:
                st.success(f"步骤1完成: 成功转换 {len(sac_files)} 个SAC文件为SAC格式")
            
            # 显示结果
            st.subheader("生成的SAC文件")
            sac_files_list = list(sac_dir.glob("**/*.sac"))
            if sac_files_list:
                sac_df = pd.DataFrame({"文件名": [f.name for f in sac_files_list], "路径": [str(f) for f in sac_files_list]})
                st.dataframe(sac_df)
            else:
                st.write("未找到SAC文件")
        
        except Exception as e:
            st.error(f"步骤1执行失败：{e}")
    
    # 步骤2: SAC时间同步
    if run_step2 and step2:
        try:
            base_dir = Path(base_output_dir)
            workflow_dir, sac_dir, sync_dir, segy_dir = create_workflow_dir(base_dir, workflow_name)
            
            with st.spinner("SAC文件时间同步..."):
                sac_files_to_sync = list_sac_files(sac_dir, recursive=True)
                if sac_files_to_sync:
                    synced_files = batch_synchronize_sac(
                        sac_files_to_sync, 
                        sync_dir,
                        target_sampling_rate=target_sampling_rate_step
                    )
                    st.success(f"步骤2完成: 成功同步 {len(synced_files)} 个SAC文件")
                    
                    # 显示结果
                    st.subheader("同步后的SAC文件")
                    synced_files_list = list(sync_dir.glob("*.sac"))
                    if synced_files_list:
                        sync_df = pd.DataFrame({"文件名": [f.name for f in synced_files_list], "路径": [str(f) for f in synced_files_list]})
                        st.dataframe(sync_df)
                    else:
                        st.write("未找到同步后的SAC文件")
                else:
                    st.warning("在SAC目录中未找到需要同步的文件。")

        except Exception as e:
            st.error(f"步骤2执行失败：{e}")
    
    # 步骤3: SAC到SEGY
    if run_step3 and step3:
        try:
            base_dir = Path(base_output_dir)
            workflow_dir, sac_dir, sync_dir, segy_dir = create_workflow_dir(base_dir, workflow_name)
            
            with st.spinner("转换SAC到SEGY..."):
                # 优先使用同步后的文件
                sac_files_to_process = list_sac_files(sync_dir, recursive=False)
                source_dir = sync_dir
                if not sac_files_to_process:
                    st.warning("未找到同步后的SAC文件，将尝试使用原始SAC文件。")
                    sac_files_to_process = list_sac_files(sac_dir, recursive=True)
                    source_dir = sac_dir

                if not sac_files_to_process:
                    st.error("未找到任何SAC文件进行转换。")
                    segy_files = []
                else:
                    if segy_mode == "每个SAC单独转换":
                        segy_files = convert_directory_sac_to_segy(
                            source_dir, 
                            segy_dir, 
                            recursive=True,
                            overwrite=True
                        )
                    else: # 合并为单个SEGY
                        output_filename = f"combined_{workflow_name}.segy"
                        output_file = segy_dir / output_filename
                        
                        files_to_merge = sac_files_to_process
                        if 'merge_z_only_step' in locals() and merge_z_only_step:
                            files_to_merge = [f for f in sac_files_to_process if f.stem.upper().endswith('_Z')]

                        if not files_to_merge:
                            st.warning("没有符合条件的SAC文件可供合并。")
                            segy_files = []
                        else:
                            trace_mapping = combine_sacs_to_single_segy(
                                files_to_merge,
                                output_file,
                                overwrite=True,
                                strict_one_trace_per_sac=strict_one,
                                split_output_files=split_output_files,
                                merge_z_only=False # UI已处理
                            )
                            
                            if split_output_files:
                                segy_files = list(segy_dir.glob(f"{output_file.stem}_part*.segy"))
                            else:
                                segy_files = [output_file]
            
            st.success(f"步骤3完成: 成功生成 {len(segy_files)} 个SEGY文件")
            
            # 显示结果
            st.subheader("生成的SEGY文件")
            if segy_files:
                segy_df = pd.DataFrame({"文件名": [Path(f).name for f in segy_files], "路径": [str(f) for f in segy_files]})
                st.dataframe(segy_df)

                if segy_mode == "合并为单个SEGY" and 'trace_mapping' in locals() and trace_mapping:
                    st.subheader("文件到道的映射信息")
                    mapping_df = pd.DataFrame({
                        "道号": [i+1 for i in range(len(trace_mapping))],
                        "原始SAC文件名": trace_mapping
                    })
                    st.dataframe(mapping_df)
            else:
                st.write("未找到SEGY文件")
        
        except Exception as e:
            st.error(f"步骤3执行失败：{e}")


# EPS2到SAC模式
elif workflow_mode == "EPS2→SAC":
    st.header("EPS2到SAC转换")
    
    run_eps2_to_sac = st.button("执行EPS2到SAC转换")
    
    if run_eps2_to_sac:
        try:
            input_dir = Path(eps2_input_dir)
            output_dir = Path(sac_output_dir)
            
            if not input_dir.is_dir():
                st.error(f"输入目录不存在: {input_dir}")
            else:
                with st.spinner("转换EPS2到SAC..."):
                    eps2_files = list_eps2_files(input_dir, recursive=recursive)
                    if not eps2_files:
                        st.error(f"在目录中未找到EPS2文件: {input_dir}")
                    else:
                        sac_files = []
                        st.info(f"找到 {len(eps2_files)} 个EPS2文件，开始转换...")
                        for eps2_file in eps2_files:
                            try:
                                files = convert_eps2_to_sac(eps2_file, output_dir)
                                sac_files.extend(files)
                                st.success(f"✓ 成功转换 {eps2_file.name} -> {len(files)} 个SAC文件")
                            except Exception as e:
                                st.error(f"✗ 转换失败 {eps2_file.name}: {e}")
                
                st.success(f"转换完成: 成功转换 {len(sac_files)} 个EPS2文件为SAC格式")
                
                # 显示结果
                st.subheader("生成的SAC文件")
                sac_files_list = list(output_dir.glob("**/*.sac"))
                if sac_files_list:
                    sac_df = pd.DataFrame({"文件名": [f.name for f in sac_files_list], "路径": [str(f) for f in sac_files_list]})
                    st.dataframe(sac_df)
                else:
                    st.write("未找到SAC文件")
                
                # 提供打开输出目录按钮
                if st.button("打开输出目录"):
                    try:
                        os.startfile(str(output_dir))
                    except Exception as e:
                        st.warning(f"无法打开目录：{e}")
        
        except Exception as e:
            st.error(f"转换失败：{e}")


# SAC时间同步模式
elif workflow_mode == "SAC时间同步":
    st.header("SAC文件时间同步")
    
    run_sync = st.button("执行SAC时间同步")
    
    if run_sync:
        try:
            input_dir = Path(sac_input_dir)
            output_dir = Path(sac_sync_output_dir)
            
            if not input_dir.is_dir():
                st.error(f"输入目录不存在: {input_dir}")
            else:
                with st.spinner("同步SAC文件..."):
                    output_dir.mkdir(parents=True, exist_ok=True)
                    sac_files_to_sync = list_sac_files(input_dir, recursive=True)
                    if not sac_files_to_sync:
                        st.warning("在输入目录中未找到SAC文件。")
                    else:
                        synced_files = batch_synchronize_sac(
                            sac_files_to_sync, 
                            output_dir,
                            target_sampling_rate=target_sampling_rate_sync
                        )
                        st.success(f"SAC文件时间同步完成，同步了 {len(synced_files)} 个文件。")
                
                # 显示结果
                st.subheader("同步后的SAC文件")
                sac_files_list = list(output_dir.glob("*.sac"))
                if sac_files_list:
                    sac_df = pd.DataFrame({"文件名": [f.name for f in sac_files_list], "路径": [str(f) for f in sac_files_list]})
                    st.dataframe(sac_df)
                else:
                    st.write("未找到同步后的SAC文件")
                
                # 提供打开输出目录按钮
                if st.button("打开输出目录"):
                    try:
                        os.startfile(str(output_dir))
                    except Exception as e:
                        st.warning(f"无法打开目录：{e}")
        
        except Exception as e:
            st.error(f"同步失败：{e}")


# SAC到SEGY模式
elif workflow_mode == "SAC→SEGY":
    st.header("SAC到SEGY转换")
    
    run_sac_to_segy = st.button("执行SAC到SEGY转换")
    
    if run_sac_to_segy:
        try:
            input_dir = Path(sac_segy_input_dir)
            output_dir = Path(segy_output_dir)
            
            if not input_dir.is_dir():
                st.error(f"输入目录不存在: {input_dir}")
            else:
                output_dir.mkdir(parents=True, exist_ok=True)
                
                with st.spinner("转换SAC到SEGY..."):
                    if segy_mode == "每个SAC单独转换":
                        # 每个SAC文件单独转换
                        segy_files = convert_directory_sac_to_segy(
                            input_dir, 
                            output_dir, 
                            recursive=True,
                            overwrite=True
                        )
                    else:
                        # 合并为单个SEGY文件
                        sac_files = list_sac_files(input_dir, recursive=True)
                        if not sac_files:
                            st.warning("输入目录中没有找到SAC文件。")
                            segy_files = []
                        else:
                            output_file = output_dir / "combined.segy"
                            
                            files_to_merge = sac_files
                            if merge_z_only_sac:
                                files_to_merge = [f for f in sac_files if f.stem.upper().endswith('_Z')]

                            if not files_to_merge:
                                st.warning("没有符合条件的SAC文件用于合并。")
                                segy_files = []
                            else:
                                trace_mapping = combine_sacs_to_single_segy(
                                    files_to_merge,
                                    output_file,
                                    overwrite=True,
                                    strict_one_trace_per_sac=strict_one,
                                    split_output_files=split_output_files,
                                    merge_z_only=False # UI已处理
                                )
                                
                                if split_output_files:
                                    # 收集分割后的文件
                                    segy_files = list(output_dir.glob("combined_part*.segy"))
                                else:
                                    segy_files = [output_file]
                
                st.success(f"转换完成: 成功生成 {len(segy_files)} 个SEGY文件")
                
                # 显示结果
                st.subheader("生成的SEGY文件")
                if segy_files:
                    segy_df = pd.DataFrame({"文件名": [Path(f).name for f in segy_files], "路径": [str(f) for f in segy_files]})
                    st.dataframe(segy_df)

                    if segy_mode == "合并为单个SEGY" and 'trace_mapping' in locals() and trace_mapping:
                        st.subheader("文件到道的映射信息")
                        mapping_df = pd.DataFrame({
                            "道号": [i+1 for i in range(len(trace_mapping))],
                            "原始SAC文件名": trace_mapping
                        })
                        st.dataframe(mapping_df)
                else:
                    st.write("未找到SEGY文件")
                
                # 提供打开输出目录按钮
                if st.button("打开输出目录"):
                    try:
                        os.startfile(str(output_dir))
                    except Exception as e:
                        st.warning(f"无法打开目录：{e}")
        
        except Exception as e:
            st.error(f"转换失败：{e}")


# 验证完整流程模式
elif workflow_mode == "验证完整流程模式":
    st.header("验证完整流程模式 (EPS2→SAC→同步→提取Z分量→合并分割SEGY)")
    st.write("此模式将执行完整的验证流程：")
    st.write("1. 读取EPS2文件并转换为SAC格式")
    st.write("2. 同步所有SAC文件的时间")
    st.write("3. 提取Z分量SAC文件")
    st.write("4. 合并Z分量SAC文件为SEGY格式")
    st.write("5. 当采样点数超过32767时自动分割文件")
    
    col1, col2 = st.columns([3, 1])
    with col1:
        eps2_dir = st.text_input("EPS2文件目录", value="EPS2data", key="eps2_dir_verification")
    with col2:
        if st.button("浏览目录", key="browse_eps2_verification"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.eps2_dir_verification = selected_dir
    
    col3, col4 = st.columns([3, 1])
    with col3:
        output_base_dir = st.text_input("输出基础目录", value="verification_test", key="output_base_dir_verification")
    with col4:
        if st.button("浏览目录", key="browse_output_verification"):
            selected_dir = select_directory()
            if selected_dir:
                st.session_state.output_base_dir_verification = selected_dir
    
    # 从session_state获取目录值
    if 'eps2_dir_verification' in st.session_state:
        eps2_dir = st.session_state.eps2_dir_verification
    if 'output_base_dir_verification' in st.session_state:
        output_base_dir = st.session_state.output_base_dir_verification
    
    target_sampling_rate_verify = st.number_input("目标采样率 (Hz)", value=100.0, min_value=0.0, format="%.2f", key="sampling_rate_verify", help="指定重采样时的目标采样率。如果设为0，则自动使用最低采样率。")

    if st.button("执行验证完整流程"):
        try:
            eps2_path = Path(eps2_dir)
            output_path = Path(output_base_dir)
            
            if not eps2_path.is_dir():
                st.error(f"EPS2目录不存在: {eps2_path}")
            else:
                # 步骤1：读取EPS2文件并转换为SAC格式
                with st.spinner("步骤1：转换EPS2文件为SAC格式..."):
                    sac_dir = output_path / "sac"
                    sac_dir.mkdir(parents=True, exist_ok=True)
                    eps2_files = list_eps2_files(eps2_path, recursive=True)
                    if not eps2_files:
                        st.error(f"在目录中未找到EPS2文件: {eps2_path}")
                        sac_files = []
                    else:
                        sac_files = []
                        failed_files = []
                        st.info(f"找到 {len(eps2_files)} 个EPS2文件，开始转换...")
                        for eps2_file in eps2_files:
                            try:
                                files = convert_eps2_to_sac(eps2_file, sac_dir)
                                sac_files.extend(files)
                                st.success(f"✓ 成功转换 {eps2_file.name} -> {len(files)} 个SAC文件")
                            except Exception as e:
                                failed_files.append((eps2_file, str(e)))
                                st.error(f"✗ 转换失败 {eps2_file.name}: {e}")
                        
                        if sac_files:
                            st.success(f"步骤1完成：成功转换了 {len(sac_files)} 个SAC文件")
                        if failed_files:
                            st.warning(f"有 {len(failed_files)} 个EPS2文件转换失败")
                
                # 步骤2：同步SAC文件
                synced_files = []
                if sac_files:
                    with st.spinner("步骤2：同步SAC文件..."):
                        sync_dir = output_path / "sync"
                        sync_dir.mkdir(parents=True, exist_ok=True)
                        
                        synced_files = batch_synchronize_sac(
                            sac_files, 
                            sync_dir,
                            target_sampling_rate=target_sampling_rate_verify
                        )
                        st.success(f"步骤2完成：同步了 {len(synced_files)} 个SAC文件")
                
                # 步骤3：提取Z分量
                z_files = []
                if synced_files:
                    with st.spinner("步骤3：提取Z分量..."):
                        z_dir = output_path / "z_component"
                        z_dir.mkdir(parents=True, exist_ok=True)
                        
                        for sac_file in synced_files:
                            sac_file_path = Path(sac_file)
                            if sac_file_path.name.upper().endswith('_Z.SAC'):
                                z_file = z_dir / sac_file_path.name
                                shutil.copy2(sac_file_path, z_file)
                                z_files.append(z_file)
                        
                        st.success(f"步骤3完成：提取了 {len(z_files)} 个Z分量文件")
                
                # 步骤4和5：合并Z分量SAC文件为SEGY格式，并在需要时分割
                segy_files = []
                trace_mapping = []
                if z_files:
                    with st.spinner("步骤4和5：合并Z分量SAC文件为SEGY格式..."):
                        segy_dir = output_path / "segy"
                        segy_dir.mkdir(parents=True, exist_ok=True)
                        
                        # 检查是否需要分割文件
                        split_files = False
                        for z_file in z_files:
                            st_temp = read(str(z_file))
                            if st_temp and st_temp[0].stats.npts > 32767:
                                split_files = True
                                break
                        
                        output_file = segy_dir / "combined_Z.segy"
                        trace_mapping = combine_sacs_to_single_segy(
                            z_files,
                            output_file,
                            overwrite=True,
                            strict_one_trace_per_sac=False,
                            split_output_files=split_files
                        )
                        
                        if split_files:
                            segy_files = list(segy_dir.glob("combined_Z_part*.segy"))
                            st.success(f"步骤4和5完成：由于采样点数超过32767，已将Z分量合并并分割为 {len(segy_files)} 个SEGY文件")
                        else:
                            segy_files = [output_file]
                            st.success(f"步骤4和5完成：已将Z分量合并为 {len(segy_files)} 个SEGY文件")
                
                # 显示SEGY文件信息
                if segy_files:
                    st.subheader("生成的SEGY文件信息")
                    segy_info = []
                    for segy_file in segy_files:
                        try:
                            st_obj = read(str(segy_file), format='SEGY')
                            if st_obj:
                                segy_info.append({
                                    "文件名": Path(segy_file).name,
                                    "道数": len(st_obj),
                                    "采样点数": st_obj[0].stats.npts,
                                    "采样率(Hz)": st_obj[0].stats.sampling_rate,
                                    "开始时间": str(st_obj[0].stats.starttime),
                                    "结束时间": str(st_obj[0].stats.endtime)
                                })
                        except Exception as read_e:
                            st.warning(f"读取SEGY文件失败: {segy_file.name}, {read_e}")

                    
                    if segy_info:
                        df = pd.DataFrame(segy_info)
                        st.dataframe(df)

                    if trace_mapping:
                        st.subheader("文件到道的映射信息")
                        mapping_df = pd.DataFrame({
                            "道号": [i+1 for i in range(len(trace_mapping))],
                            "原始SAC文件名": trace_mapping
                        })
                        st.dataframe(mapping_df)
                
                # 显示输出目录按钮
                if st.button("打开输出目录"):
                    try:
                        os.startfile(str(output_path))
                    except Exception as e:
                        st.warning(f"无法打开目录：{e}")
                        
        except Exception as e:
            st.error(f"验证流程失败: {e}")


st.divider()
st.caption(
    "提示：本工具支持完整的EPS2到SEGY转换流程，包括EPS2→SAC转换、SAC文件时间同步和SAC→SEGY转换。"
)
